#include "utils.h"

#include "modular_reduction.h"
using namespace std;

/*
 * @brief 计算模幂运算 \( x^e \mod p \)。
 * @param x 被乘数，底数，64 位无符号整数。
 * @param e 指数，64 位无符号整数。
 * @param p 模数，64 位无符号整数，必须大于 0。
 * @return 计算结果 \( x^e \mod p \)，64 位无符号整数。
 */
uint64_t ModExp(uint64_t x, uint64_t e, uint64_t p)
{
    vector<uint64_t> params = BRedParams(p);
    uint64_t result = 1;
    while (e > 0) {
        if (e & 1) {
            result = BRed(result, x, p, params);
        }
        x = BRed(x, x, p, params);
        e >>= 1;
    }
    return result;
}
/*
 * @brief 比较两个 uint64_t 类型的向量是否相等。
 * @param a 第一个要比较的向量，类型为 `const std::vector<uint64_t>&`。
 * @param b 第二个要比较的向量，类型为 `const std::vector<uint64_t>&`。
 * @return
 * - `true`：如果两个向量具有相同的大小并且所有对应元素都相等。
 * - `false`：如果两个向量的大小不同，或者至少有一个对应元素不相等。
 */
bool EqualSliceUint64(const vector<uint64_t> &a, const vector<uint64_t> &b)
{
    if (a.size() != b.size())
        return false;
    return equal(a.begin(), a.end(), b.begin());
}

/*
 * @brief 反转一个 64 位无符号整数的前 `bitLen` 位。
 * @param x 要反转比特位的 64 位无符号整数。
 * @param bitLen 需要反转的比特位长度，必须满足 `1 <= bitLen <= 64`。
 * @return 反转后的结果，仅保留前 `bitLen` 位的反转值，类型为 `uint64_t`。
 */
uint64_t BitReverse64(uint64_t x, uint64_t bitLen)
{
    x = ((x & 0x5555555555555555ULL) << 1) | ((x >> 1) & 0x5555555555555555ULL);
    x = ((x & 0x3333333333333333ULL) << 2) | ((x >> 2) & 0x3333333333333333ULL);
    x = ((x & 0x0F0F0F0F0F0F0F0FULL) << 4) | ((x >> 4) & 0x0F0F0F0F0F0F0F0FULL);
    x = ((x & 0x00FF00FF00FF00FFULL) << 8) | ((x >> 8) & 0x00FF00FF00FF00FFULL);
    x = ((x & 0x0000FFFF0000FFFFULL) << 16) | ((x >> 16) & 0x0000FFFF0000FFFFULL);
    x = (x << 32) | (x >> 32);

    return x >> (64 - bitLen);
}
/*
 * @brief 检查向量中的所有元素是否互不相同。
 *
 * 该函数遍历输入的 `std::vector<uint64_t>`，并使用 `std::unordered_set` 来判断所有元素是否唯一。
 *
 * @param s 要检查的 `std::vector<uint64_t>`。
 * @return 如果所有元素都不同，返回 `true`；否则，返回 `false`。
 */
bool AllDistinct(const vector<uint64_t> &s)
{
    unordered_set<uint64_t> set;
    for (uint64_t si : s) {
        if (!set.insert(si).second) {
            return false;
        }
    }
    return true;
}
/*
 * @brief 计算一个 64 位无符号整数的二进制位长度。该函数通过不断右移位来计算整数 `x` 的二进制表示中有效位的数量。
 * @param x 要计算位长度的 64 位无符号整数。
 * @return `x` 的二进制位长度。如果 `x` 为 0，返回 0。
 */
int BitsLen64(uint64_t x)
{
    if (x == 0)
        return 0;
    int len = 0;
    while (x != 0) {
        len++;
        x >>= 1;
    }
    return len;
}

/*
 * @brief 判断一个 64 位无符号整数是否为素数。该函数使用 GMP 库的 `mpz_probab_prime_p`
 * 函数进行素性测试，具有一定的概率性。
 * @param x 要判断的 64 位无符号整数。
 * @return 如果 `x` 是素数，返回 `true`；否则，返回 `false`。
 */
bool IsPrime(uint64_t x)
{
    mpz_class num(static_cast<unsigned long int>(x));
    int is_prime = mpz_probab_prime_p(num.get_mpz_t(), 25);
    return is_prime != 0;
}
/*
 * @brief 生成一组素数Qmul，为了在新的模数下进行密文乘
 * @param logP 新模数的bit数。
 * @param NthRoot 2*N，N为密文向量的长度。
 * @param n 新模数的长度
 * @return 长度为n的新模数vector容器。
 */
vector<uint64_t> GenerateNTTPrimesP(uint64_t logP, uint64_t NthRoot, uint64_t n)
{
    vector<uint64_t> primes;
    uint64_t Ppow2 = 1ULL << logP;
    uint64_t x = Ppow2 + 1;
    while (true) {
        if (x > NthRoot) {
            x -= NthRoot;
            if (IsPrime(x)) {
                primes.push_back(x);
                if (primes.size() == n) {
                    return primes;
                }
            }
        } else {
            throw runtime_error("generateNTTPrimesP error: cannot generate enough primes for the given parameters");
        }
    }
}

/*
 * @brief 创建一个新的 `mpz_class` 类型，并将给定的无符号长整型值转换为大整数。
 * 该函数使用 GMP 库的 `mpz_set_ui` 将输入的无符号长整型值 `v` 转换为 `mpz_class` 类型。
 * @param v 要转换的无符号长整型值。
 * @return 转换后的 `mpz_class` 类型的大整数。
 */
mpz_class NewUint(unsigned long int v)
{
    mpz_class result;
    mpz_set_ui(result.get_mpz_t(), static_cast<unsigned long int>(v));
    return result;
}

/*
 * @brief 返回给定 `uint64_t` 切片中的最大值。
 * @param slice 要查找最大值的 `vector<uint64_t>`。
 * @param len 切片的长度，即 `slice` 向量的元素数量。
 * @return 切片中的最大值。
 */
uint64_t MaxSliceUint64(vector<uint64_t> slice, int len)
{
    uint64_t res = slice[0];
    for (uint64_t i = 1; i < len; i++) {
        res = max(res, slice[i]);
    }
    return res;
}

/*
 * @brief 将 `uint64_t` 类型的切片（向量）旋转 `k` 个位置。
 * 该函数将输入的 `vector<uint64_t>` 向量分成两部分，并分别旋转这两部分，然后合并旋转后的结果。
 * @param s 输入的 `vector<uint64_t>` 向量。
 * @param k 旋转的步数（正数表示右旋，负数表示左旋）。
 * @return 旋转后的新向量。
 */
vector<uint64_t> RotateUint64Slots(vector<uint64_t> s, int64_t k)
{
    vector<uint64_t> ret(s.size());
    uint64_t slots = s.size() >> 1;

    // 旋转前半部分
    vector<uint64_t> tmp = RotateUint64Slice(vector<uint64_t>(s.begin(), s.begin() + slots), k);
    copy(tmp.begin(), tmp.end(), ret.begin());

    tmp.clear();

    // 旋转后半部分
    tmp = RotateUint64Slice(vector<uint64_t>(s.begin() + slots, s.end()), k);
    copy(tmp.begin(), tmp.end(), ret.begin() + slots);

    return ret;
}
/*
 * @brief 将给定的 `uint64_t` 向量左移 `k` 个位置,k为负数就相当于右移-k位。
 * @param s 输入的 `vector<uint64_t>` 向量。
 * @param k 左移的步数，正数表示右移，负数表示左移。
 * @return 旋转后的新向量。
 */
vector<uint64_t> RotateUint64Slice(vector<uint64_t> s, int64_t k)
{
    if (k == 0 || s.size() == 0) {
        return s;
    }
    int64_t r = k % s.size();
    if (r < 0) {
        r = r + s.size();
    }
    vector<uint64_t> ret(s.size());

    copy(s.begin() + r, s.end(), ret.begin());
    copy(s.begin(), s.begin() + r, ret.begin() + s.size() - r);
    return ret;
}
/*
 * @brief 计算 (operand1 - operand2) % modulus，并调用加速库的减法函数。
 * @param operand1 被减数。
 * @param operand2 减数。
 * @param modulus 模数。
 * @return 计算结果 (operand1 - operand2) % modulus。
 */
uint64_t SubUintMod(uint64_t operand1, uint64_t operand2, const uint64_t &modulus)
{
    uint64_t result;
    SubMod(result, operand1, operand2, modulus); // 调用加速库中的SubMod函数
    return result;
}

/*
 * @brief 计算 (operand1 + operand2) % modulus，并调用加速库的加法函数。
 * @param operand1 加数。
 * @param operand2 加数。
 * @param modulus 模数。
 * @return 计算结果 (operand1 + operand2) % modulus。
 */
uint64_t AddUintMod(uint64_t operand1, uint64_t operand2, const uint64_t &modulus)
{
    uint64_t result = 0;
    AddMod(result, operand1, operand2, modulus); // 调用加速库中的AddMod函数
    return result;
}

/*
 * @brief 计算 (operand1 * operand2) % modulus，并调用加速库的乘法函数。
 * @param operand1 被乘数。
 * @param operand2 乘数。
 * @param modulus 模数。
 * @return 计算结果 (operand1 * operand2) % modulus。
 */
uint64_t MultiplyUintMod(uint64_t operand1, uint64_t operand2, const uint64_t &modulus)
{
    uint64_t result = 0;
    MulMod(result, operand1, operand2, modulus); // 调用加速库中的MulMod函数
    return result;
}

/*
 * @brief 计算 -base % modulus。
 * @param base 操作数。
 * @param modulus 模数。
 * @return 计算结果 -base % modulus。
 */
uint64_t NegateUintMod(uint64_t base, const uint64_t &modulus)
{
    if (modulus == 0) {
        cout << "modulus模数为0 " << endl;
    }
    // 确保 base 小于 modulus
    while (base > modulus) {
        base -= modulus;
    }
    uint64_t result = (modulus - base) % modulus; // 计算负数的模
    return result;
}

/*
 * @brief 计算 base^exp
 * @param base 底数，即需要进行指数运算的数。
 * @param exp 指数，即幂次。
 * @return 计算结果 base^exponent。
 */
uint64_t Exponentiate(uint64_t base, uint64_t exp)
{
    if (exp == 0) {
        return 1;
    }
    if (exp == 1) {
        return base;
    }
    uint64_t prodult = 0;
    uint64_t res = 1;
    uint64_t currentPower = base;
    while (true) {
        if (exp & 1) {
            prodult = currentPower * res;
            swap(prodult, res);
        }
        exp >>= 1;
        if (exp == 0) {
            break;
        }
        prodult = currentPower * currentPower;
        swap(prodult, currentPower);
    }
    return res;
}